package com.abc.config;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * Milvus向量存储配置类
 * 
 * @author abc
 * @since 2024-01-01
 */
@Slf4j
@Configuration
public class VectorStoreConfig {

    @Value("${spring.ai.vectorstore.milvus.client.host:localhost}")
    private String host;

    @Value("${spring.ai.vectorstore.milvus.client.port:19530}")
    private int port;

    @Value("${spring.ai.vectorstore.milvus.client.username:root}")
    private String username;

    @Value("${spring.ai.vectorstore.milvus.client.password:milvus}")
    private String password;

    @Value("${spring.ai.vectorstore.milvus.databaseName:knowledge_platform}")
    private String databaseName;

    @Value("${spring.ai.vectorstore.milvus.collectionName:embedding}")
    private String collectionName;

    @Value("${spring.ai.vectorstore.milvus.embeddingDimension:1024}")
    private int embeddingDimension;

    @Value("${spring.ai.vectorstore.milvus.indexType:IVF_FLAT}")
    private String indexType;

    @Value("${spring.ai.vectorstore.milvus.metricType:COSINE}")
    private String metricType;

    @Autowired
    private EmbeddingModel alibabaCloudEmbeddingModel;

    /**
     * 创建Milvus客户端
     */
    @Bean
    public MilvusServiceClient milvusClient() {
        try {
            log.info("初始化Milvus客户端，连接地址: {}:{}", host, port);
            
            MilvusServiceClient client = new MilvusServiceClient(
                ConnectParam.newBuilder()
                    .withHost(host)
                    .withPort(port)
                    .withAuthorization(username, password)
                    .build()
            );
            
            log.info("Milvus客户端初始化成功");
            return client;
        } catch (Exception e) {
            log.error("Milvus客户端初始化失败", e);
            throw new RuntimeException("Milvus客户端初始化失败", e);
        }
    }

    /**
     * 创建向量存储
     */
    @Bean
    public VectorStore milvusVectorStore(MilvusServiceClient milvusClient) {
        try {
            log.info("初始化Milvus向量存储，数据库: {}, 集合: {}", databaseName, collectionName);
            
            IndexType indexTypeEnum = IndexType.valueOf(indexType);
            MetricType metricTypeEnum = MetricType.valueOf(metricType);
            
            VectorStore vectorStore = MilvusVectorStore.builder(milvusClient, alibabaCloudEmbeddingModel)
                    .collectionName(collectionName)
                    .databaseName(databaseName)
                    .indexType(indexTypeEnum)
                    .metricType(metricTypeEnum)
                    .embeddingDimension(embeddingDimension)
                    .batchingStrategy(new TokenCountBatchingStrategy())
                    .initializeSchema(true)
                    .build();
            
            log.info("Milvus向量存储初始化成功");
            return vectorStore;
        } catch (Exception e) {
            log.error("Milvus向量存储初始化失败", e);
            throw new RuntimeException("Milvus向量存储初始化失败", e);
        }
    }
}
